#################################
#   用法说明
#################################
# python color_correction.py --reference reference.jpg --input examples/03.jpg

#################################
#   导入依赖库
#################################
# 导入所需的包
from imutils.perspective import four_point_transform
from skimage import exposure
import numpy as np
import argparse
import imutils
import cv2
import sys

#################################
#   函数定义
#################################
def find_color_card(image):
    # 加载 ArUCo 字典，获取参数，并在输入图像中检测标记
    arucoDict = cv2.aruco.getPredefinedDictionary(cv2.aruco.DICT_ARUCO_ORIGINAL)
    arucoParams = cv2.aruco.DetectorParameters()
    detector = cv2.aruco.ArucoDetector(arucoDict, arucoParams)
    (corners, ids, rejected) = detector.detectMarkers(image)
    # 尝试提取色卡的四个角点坐标
    try:
        # 如果找到了四个 ArUCo 标记，则继续处理，先将ID列表展平
        ids = ids.flatten()
        # 提取左上角标记
        i = np.squeeze(np.where(ids == 923))
        topLeft = np.squeeze(corners[i])[0]
        # 提取右上角标记
        i = np.squeeze(np.where(ids == 1001))
        topRight = np.squeeze(corners[i])[1]
        # 提取右下角标记
        i = np.squeeze(np.where(ids == 241))
        bottomRight = np.squeeze(corners[i])[2]
        # 提取左下角标记
        i = np.squeeze(np.where(ids == 1007))
        bottomLeft = np.squeeze(corners[i])[3]
    # 如果未找到色卡，则直接返回 None
    except:
        return None
    # 构建参考点列表，并做透视变换，获得色卡的俯视图
    cardCoords = np.array([topLeft, topRight, bottomRight, bottomLeft])
    card = four_point_transform(image, cardCoords)
    # 返回色卡图像
    return card


# 构建参数解析器并解析命令行参数
# ap = argparse.ArgumentParser()
# ap.add_argument("-r", "--reference", required=True, help="参考图像的路径")
# ap.add_argument("-i", "--input", required=True, help="待色彩校正的输入图像路径")
# args = vars(ap.parse_args())

# 从磁盘加载参考图像和输入图像
print("[信息] 正在加载图像...")
ref = cv2.imread(args["reference"])
image = cv2.imread(args["input"])

# 调整参考图像和输入图像的大小
ref = imutils.resize(ref, width=600)
image = imutils.resize(image, width=600)

# 显示参考图像和输入图像
cv2.imshow("Reference", ref)
cv2.imshow("Input", image)

# 查找每张图像中的色卡
print("[信息] 正在查找色卡...")
refCard = find_color_card(ref)
imageCard = find_color_card(image)

# 如果在参考图像或输入图像中未找到色卡，则退出程序
if refCard is None or imageCard is None:
    print("[信息] 未能在两张图像中都找到色卡！程序退出...")
    sys.exit(0)

# 分别显示参考图像和输入图像中的色卡
cv2.imshow("Reference Color Card", refCard)
cv2.imshow("Input Color Card", imageCard)

# 将参考图像色卡的直方图匹配到输入图像色卡，实现色彩校正
print("[信息] 正在进行直方图匹配...")
imageCard = exposure.match_histograms(imageCard, refCard, channel_axis=-1)

# 显示直方图匹配后的输入色卡
cv2.imshow("Input Color Card After Matching", imageCard)
cv2.waitKey(0)

